"""
Dify API 处理器
处理ComfyUI与Dify之间的API通信和数据格式转换
"""

import json
import logging
from typing import Dict, Any, List, Optional, Union
from datetime import datetime

logger = logging.getLogger(__name__)

class DifyAPIHandler:
    """
    Dify API 处理器类
    负责处理ComfyUI输出格式到Dify期望格式的转换
    """
    
    def __init__(self):
        self.supported_formats = ['text', 'json', 'files']
        logger.info("[DifyAPIHandler] 初始化完成")
    
    def format_for_dify(self, data: Any, output_type: str = 'text') -> Dict[str, Any]:
        """
        将ComfyUI输出格式化为Dify期望的格式
        
        Args:
            data: 要格式化的数据
            output_type: 输出类型 ('text', 'json', 'files')
            
        Returns:
            Dict: Dify格式的输出
        """
        try:
            # 基础的Dify输出结构
            dify_output = {
                "text": "",
                "files": [],
                "json": [{"data": []}]
            }
            
            if output_type == 'text':
                dify_output = self._format_text_output(data, dify_output)
            elif output_type == 'json':
                dify_output = self._format_json_output(data, dify_output)
            elif output_type == 'files':
                dify_output = self._format_files_output(data, dify_output)
            else:
                logger.warning(f"不支持的输出类型: {output_type}")
                dify_output["text"] = str(data)
            
            logger.info(f"成功格式化为Dify格式，类型: {output_type}")
            return dify_output
            
        except Exception as e:
            logger.error(f"格式化数据时出错: {str(e)}")
            return {
                "text": f"格式化错误: {str(e)}",
                "files": [],
                "json": [{"data": []}]
            }
    
    def _format_text_output(self, data: Any, base_output: Dict[str, Any]) -> Dict[str, Any]:
        """格式化文本输出"""
        if isinstance(data, str):
            base_output["text"] = data
        elif isinstance(data, dict) and "text" in data:
            base_output["text"] = data["text"]
            # 如果有其他数据，放到json字段中
            if len(data) > 1:
                other_data = {k: v for k, v in data.items() if k != "text"}
                base_output["json"][0]["data"].append(other_data)
        else:
            base_output["text"] = str(data)
        
        return base_output
    
    def _format_json_output(self, data: Any, base_output: Dict[str, Any]) -> Dict[str, Any]:
        """格式化JSON输出"""
        try:
            if isinstance(data, str):
                # 尝试解析JSON字符串
                try:
                    parsed_data = json.loads(data)
                    base_output["json"][0]["data"] = [parsed_data]
                    base_output["text"] = f"JSON数据 (长度: {len(str(parsed_data))})"
                except json.JSONDecodeError:
                    base_output["text"] = data
                    base_output["json"][0]["data"] = [{"raw_text": data}]
            elif isinstance(data, (dict, list)):
                base_output["json"][0]["data"] = [data] if isinstance(data, dict) else data
                base_output["text"] = f"结构化数据 (类型: {type(data).__name__})"
            else:
                base_output["json"][0]["data"] = [{"value": data, "type": type(data).__name__}]
                base_output["text"] = str(data)
        
        except Exception as e:
            logger.error(f"格式化JSON数据时出错: {str(e)}")
            base_output["text"] = f"JSON格式化错误: {str(e)}"
            base_output["json"][0]["data"] = [{"error": str(e)}]
        
        return base_output
    
    def _format_files_output(self, data: Any, base_output: Dict[str, Any]) -> Dict[str, Any]:
        """格式化文件输出"""
        # TODO: 实现文件输出格式化
        # 这里可以处理图片、音频等文件类型
        base_output["text"] = "文件输出功能待实现"
        base_output["files"] = []
        
        return base_output
    
    def process_comfyui_result(self, result: Dict[str, Any]) -> Dict[str, Any]:
        """
        处理ComfyUI工作流的完整结果
        
        Args:
            result: ComfyUI工作流执行结果
            
        Returns:
            Dict: 处理后的Dify格式结果
        """
        try:
            # 检查结果中是否有文本输出节点的结果
            text_outputs = []
            json_outputs = []
            
            for node_id, node_result in result.items():
                if isinstance(node_result, dict):
                    # 查找文本输出
                    if "text" in node_result:
                        text_outputs.append(node_result["text"])
                    
                    # 查找JSON输出
                    if "json_output" in node_result:
                        json_outputs.append(node_result["json_output"])
            
            # 构建最终输出
            final_output = {
                "text": "",
                "files": [],
                "json": [{"data": []}]
            }
            
            # 合并文本输出
            if text_outputs:
                final_output["text"] = "\n".join(text_outputs)
            
            # 合并JSON输出
            if json_outputs:
                final_output["json"][0]["data"] = json_outputs
            
            # 如果没有找到特定的输出，返回原始结果
            if not text_outputs and not json_outputs:
                final_output["text"] = "ComfyUI工作流执行完成"
                final_output["json"][0]["data"] = [result]
            
            logger.info(f"处理ComfyUI结果完成，找到 {len(text_outputs)} 个文本输出")
            return final_output
            
        except Exception as e:
            logger.error(f"处理ComfyUI结果时出错: {str(e)}")
            return {
                "text": f"处理结果时出错: {str(e)}",
                "files": [],
                "json": [{"data": [{"error": str(e), "original_result": result}]}]
            }
    
    def validate_dify_format(self, data: Dict[str, Any]) -> bool:
        """
        验证数据是否符合Dify期望的格式
        
        Args:
            data: 要验证的数据
            
        Returns:
            bool: 是否符合格式
        """
        required_keys = ["text", "files", "json"]
        
        if not isinstance(data, dict):
            logger.warning("数据不是字典格式")
            return False
        
        for key in required_keys:
            if key not in data:
                logger.warning(f"缺少必需的键: {key}")
                return False
        
        # 验证json字段的结构
        if not isinstance(data["json"], list) or len(data["json"]) == 0:
            logger.warning("json字段格式不正确")
            return False
        
        if not isinstance(data["json"][0], dict) or "data" not in data["json"][0]:
            logger.warning("json[0]缺少data字段")
            return False
        
        logger.info("Dify格式验证通过")
        return True
    
    def create_error_response(self, error_message: str, original_data: Any = None) -> Dict[str, Any]:
        """
        创建错误响应
        
        Args:
            error_message: 错误信息
            original_data: 原始数据（可选）
            
        Returns:
            Dict: 错误响应
        """
        error_response = {
            "text": f"错误: {error_message}",
            "files": [],
            "json": [{
                "data": [{
                    "error": error_message,
                    "timestamp": datetime.now().isoformat(),
                    "original_data": original_data
                }]
            }]
        }
        
        logger.error(f"创建错误响应: {error_message}")
        return error_response

# 全局API处理器实例
api_handler = DifyAPIHandler()

def format_for_dify(data: Any, output_type: str = 'text') -> Dict[str, Any]:
    """
    便捷函数：格式化数据为Dify格式
    """
    return api_handler.format_for_dify(data, output_type)

def process_comfyui_result(result: Dict[str, Any]) -> Dict[str, Any]:
    """
    便捷函数：处理ComfyUI结果
    """
    return api_handler.process_comfyui_result(result) 